#!/bin/bash

# 设置颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 输出文件
OUTPUT_FILE="tcp_connection_analysis_$(date +%Y%m%d_%H%M%S).txt"

# 默认值
INTERVAL=5      # 默认收集频率为5秒
DURATION=60     # 默认收集时长为60秒
INTERFACE=""     # 网卡接口名称，默认自动检测

# 帮助信息
show_help() {
    echo "用法: $0 [选项]"
    echo "选项:"
    echo "  -o, --output <文件>   输出文件（默认: tcp_connection_analysis_时间戳.txt）"
    echo "  -i, --interval <秒>   收集频率，每隔多少秒收集一次数据（默认: 5秒）"
    echo "  -d, --duration <秒>   收集时长，收集数据的总时间（默认: 60秒）"
    echo "  -e, --interface <名称> 指定网卡名称（默认: 自动检测默认网卡）"
    echo "  -h, --help           显示此帮助信息"
    exit 0
}

# 解析命令行参数
while [[ $# -gt 0 ]]; do
    case $1 in
        -o|--output)
            OUTPUT_FILE="$2"
            shift 2
            ;;
        -i|--interval)
            INTERVAL="$2"
            shift 2
            ;;
        -d|--duration)
            DURATION="$2"
            shift 2
            ;;
        -e|--interface)
            INTERFACE="$2"
            shift 2
            ;;
        -h|--help)
            show_help
            ;;
        *)
            echo "未知选项: $1"
            show_help
            ;;
    esac
done

# 检查参数有效性
if ! [[ "$INTERVAL" =~ ^[0-9]+$ ]]; then
    echo -e "${RED}错误: 收集频率必须是正整数${NC}"
    exit 1
fi

if ! [[ "$DURATION" =~ ^[0-9]+$ ]]; then
    echo -e "${RED}错误: 收集时长必须是整数${NC}"
    exit 1
fi

# 检查ss命令是否可用
if ! command -v ss &> /dev/null; then
    echo -e "${RED}错误: ss 命令不可用，请安装 iproute2 包${NC}"
    exit 1
fi

# 获取默认网卡
get_default_interface() {
    local default_iface=""
    
    # 尝试方法1：通过默认路由获取
    if command -v ip &> /dev/null; then
        default_iface=$(ip route | grep default | awk '{print $5}' | head -n 1)
    fi
    
    # 尝试方法2：如果方法1失败，通过ifconfig查找活跃网卡
    if [ -z "$default_iface" ] && command -v ifconfig &> /dev/null; then
        default_iface=$(ifconfig | grep -E "^[a-zA-Z0-9]+" | grep -v "lo" | head -n 1 | cut -d ":" -f 1)
    fi
    
    # 如果还是失败，使用回退值
    if [ -z "$default_iface" ]; then
        echo -e "${YELLOW}警告: 无法自动检测默认网卡，使用eth0${NC}" >&2
        default_iface="eth0"
    fi
    
    echo "$default_iface"
}

# 如果未指定网卡，则获取默认网卡
if [ -z "$INTERFACE" ]; then
    INTERFACE=$(get_default_interface)
fi

echo -e "${BLUE}使用网卡: $INTERFACE${NC}"

# 创建输出文件头部
echo "=== TCP连接分析报告 ===" > "$OUTPUT_FILE"
echo "生成时间: $(date)" >> "$OUTPUT_FILE"
echo "监控网卡: ${INTERFACE}" >> "$OUTPUT_FILE"
if [ "$DURATION" -gt 0 ]; then
    echo "收集频率: 每 $INTERVAL 秒" >> "$OUTPUT_FILE"
    echo "收集时长: $DURATION 秒" >> "$OUTPUT_FILE"
    echo "收集次数: $(( DURATION / INTERVAL + (DURATION % INTERVAL > 0 ? 1 : 0) )) 次" >> "$OUTPUT_FILE"
else
    echo "单次收集" >> "$OUTPUT_FILE"
fi
echo "----------------------------------------" >> "$OUTPUT_FILE"

# 临时文件用于累积多次收集的结果
temp_all_output="/tmp/ss_all_output_$$.tmp"
touch "$temp_all_output"

# 清理函数
cleanup() {
    # 清理临时文件
    rm -f "$temp_all_output"
    
    # 显示完成消息，不考虑退出参数
    echo -e "\n${GREEN}TCP连接分析完成，报告已保存到: $OUTPUT_FILE${NC}"
}

# 设置清理陷阱，但不再传递参数
trap 'cleanup' EXIT INT TERM

# 获取指定网卡的IP地址
get_interface_ip() {
    local interface=$1
    local ip_addresses=""
    
    # 获取接口的IPv4地址
    if command -v ip &> /dev/null; then
        ip_addresses=$(ip -4 addr show $interface 2>/dev/null | grep -oP 'inet \K[0-9.]+')
    elif command -v ifconfig &> /dev/null; then
        ip_addresses=$(ifconfig $interface 2>/dev/null | grep -oP 'inet \K[0-9.]+')
    fi
    
    echo "$ip_addresses"
}

# 收集数据函数
collect_data() {
    local timestamp=$(date +"%Y-%m-%d %H:%M:%S")
    
    # 获取网卡的IP地址
    local ip_addresses=$(get_interface_ip "$INTERFACE")
    local ss_filter=""
    
    # 如果能获取到IP地址，使用IP地址过滤
    if [ -n "$ip_addresses" ]; then
        # 为每个IP地址创建过滤条件
        local filter_parts=""
        for ip in $ip_addresses; do
            if [ -n "$filter_parts" ]; then
                filter_parts="$filter_parts or "
            fi
            filter_parts="${filter_parts}src $ip or dst $ip"
        done
        
        if [ -n "$filter_parts" ]; then
            ss_output=$(ss -tie "( $filter_parts )" 2>/dev/null)
        else
            # 使用网卡接口名称过滤
            ss_output=$(ss -tie dev "$INTERFACE" 2>/dev/null)
        fi
    else
        # 使用网卡接口名称过滤
        ss_output=$(ss -tie dev "$INTERFACE" 2>/dev/null)
    fi
    
    # 如果过滤后没有结果，则使用不过滤的方式并添加警告
    if [ -z "$ss_output" ]; then
        echo -e "${YELLOW}警告: 在网卡 $INTERFACE 上未找到TCP连接，显示所有TCP连接${NC}" >&2
        ss_output=$(ss -tieH)  # -H 跳过表头
    fi
    
    # 保存到临时文件
    local temp_ss_output="/tmp/ss_output_$$.tmp"
    echo "$ss_output" > "$temp_ss_output"
    
    # 追加到累积文件 - 添加时间戳标记
    echo "" >> "$temp_all_output" # 添加空行
    echo "================================================================" >> "$temp_all_output"
    echo "=== 收集时间: $timestamp ===" >> "$temp_all_output"
    echo "================================================================" >> "$temp_all_output"
    echo "" >> "$temp_all_output" # 添加空行
    cat "$temp_ss_output" >> "$temp_all_output"
    echo "" >> "$temp_all_output"
    
    # 处理当前数据
    process_data "$temp_ss_output" "$timestamp"
    
    # 清理临时文件
    rm -f "$temp_ss_output"
}

# 处理数据函数
process_data() {
    local temp_ss_output="$1"
    local timestamp="$2"
    
    # 仅在多次收集时显示时间戳
    if [ "$DURATION" -gt 0 ]; then
        echo -e "\n\n" >> "$OUTPUT_FILE" # 增加额外的空行
        echo "================================================================" >> "$OUTPUT_FILE"
        echo "=== 收集时间: $timestamp ===" >> "$OUTPUT_FILE"
        echo "================================================================" >> "$OUTPUT_FILE"
        echo -e "\n" >> "$OUTPUT_FILE" # 再增加一行空行
    fi
    
    # 统计不同状态的TCP连接
    echo -e "\n=== TCP连接状态统计 ===" >> "$OUTPUT_FILE"
    echo "----------------------------------------" >> "$OUTPUT_FILE"
    # 只获取连接信息行（不以空格开头的行），然后提取第一列（状态列）
    # 过滤掉表头（State字段）并排除空行
    connection_states=$(cat "$temp_ss_output" | grep -v "^[[:space:]]" | grep -v "^State" | awk '{print $1}' | grep -v "^$" | sort | uniq -c)
    echo "状态        连接数" >> "$OUTPUT_FILE"
    echo "----------------" >> "$OUTPUT_FILE"
    echo "$connection_states" | while read count state; do
        if [ -n "$state" ]; then
            printf "%-12s %s\n" "$state" "$count" >> "$OUTPUT_FILE"
        fi
    done

    total_connections=$(echo "$connection_states" | awk '{sum+=$1} END {print sum}')
    echo "----------------" >> "$OUTPUT_FILE"
    printf "%-12s %s\n" "总计" "$total_connections" >> "$OUTPUT_FILE"

    # 分析每个TCP连接并检测问题
    echo -e "\n=== 问题TCP连接分析 ===" >> "$OUTPUT_FILE"
    echo "----------------------------------------" >> "$OUTPUT_FILE"

    # 初始化计数器
    line_number=1
    connection_info=""
    problem_found=false
    problem_count=0
    high_rtt_count=0
    high_retrans_count=0
    window_exhausted_count=0
    rate_drop_count=0
    mtu_problem_count=0
    reordering_count=0

    # 检查高延迟问题
    found_high_rtt=false
    while IFS= read -r line; do
        if [[ "$line" =~ ^[[:space:]] ]]; then
            if [[ "$line" =~ rtt:([0-9\.]+) ]]; then
                rtt=${BASH_REMATCH[1]}
                # 简化比较，转为整数比较
                rtt_int=${rtt%%.*}
                if [ "$rtt_int" -gt 100 ]; then
                    if [ "$found_high_rtt" = false ]; then
                        echo "高延迟：（rtt:>100ms）" >> "$OUTPUT_FILE"
                        echo "" >> "$OUTPUT_FILE"
                        found_high_rtt=true
                    fi
                    echo "$connection_info" >> "$OUTPUT_FILE"
                    echo "$line" >> "$OUTPUT_FILE"
                    echo "" >> "$OUTPUT_FILE"
                    high_rtt_count=$((high_rtt_count + 1))
                    problem_count=$((problem_count + 1))
                fi
            fi
        else
            connection_info="$line"
        fi
    done < "$temp_ss_output"

    # 如果没有高延迟问题，简洁输出
    if [ "$found_high_rtt" = false ]; then
        echo "高延迟：（rtt:>100ms） - 未发现" >> "$OUTPUT_FILE"
    fi

    # 检查高重传问题
    echo "----------------------------------------" >> "$OUTPUT_FILE"
    found_high_retrans=false
    while IFS= read -r line; do
        if [[ "$line" =~ ^[[:space:]] ]]; then
            if [[ "$line" =~ retrans:|retransmit ]]; then
                if [ "$found_high_retrans" = false ]; then
                    echo "高重传：（retrans:>0 或 包含retransmit）" >> "$OUTPUT_FILE"
                    echo "" >> "$OUTPUT_FILE"
                    found_high_retrans=true
                fi
                echo "$connection_info" >> "$OUTPUT_FILE"
                echo "$line" >> "$OUTPUT_FILE"
                echo "" >> "$OUTPUT_FILE"
                high_retrans_count=$((high_retrans_count + 1))
                problem_count=$((problem_count + 1))
            fi
        else
            connection_info="$line"
        fi
    done < "$temp_ss_output"

    # 如果没有高重传问题，简洁输出
    if [ "$found_high_retrans" = false ]; then
        echo "高重传：（retrans:>0 或 包含retransmit） - 未发现" >> "$OUTPUT_FILE"
    fi

    # 检查窗口耗尽问题
    echo "----------------------------------------" >> "$OUTPUT_FILE"
    found_window_exhausted=false
    while IFS= read -r line; do
        if [[ "$line" =~ ^[[:space:]] ]]; then
            if [[ "$line" =~ cwnd:0|rcv_space:0 ]]; then
                if [ "$found_window_exhausted" = false ]; then
                    echo "窗口耗尽：（cwnd:0或rcv_space:0）" >> "$OUTPUT_FILE"
                    echo "" >> "$OUTPUT_FILE"
                    found_window_exhausted=true
                fi
                echo "$connection_info" >> "$OUTPUT_FILE"
                echo "$line" >> "$OUTPUT_FILE"
                echo "" >> "$OUTPUT_FILE"
                window_exhausted_count=$((window_exhausted_count + 1))
                problem_count=$((problem_count + 1))
            fi
        else
            connection_info="$line"
        fi
    done < "$temp_ss_output"

    # 如果没有窗口耗尽问题，简洁输出
    if [ "$found_window_exhausted" = false ]; then
        echo "窗口耗尽：（cwnd:0或rcv_space:0） - 未发现" >> "$OUTPUT_FILE"
    fi

    # 检查速率问题
    echo "----------------------------------------" >> "$OUTPUT_FILE"
    found_rate_drop=false
    while IFS= read -r line; do
        if [[ "$line" =~ ^[[:space:]] ]]; then
            if [[ "$line" =~ pacing_rate[[:space:]]0|delivery_rate[[:space:]]0 ]]; then
                if [ "$found_rate_drop" = false ]; then
                    echo "速率问题：（pacing_rate:0或delivery_rate:0）" >> "$OUTPUT_FILE"
                    echo "" >> "$OUTPUT_FILE"
                    found_rate_drop=true
                fi
                echo "$connection_info" >> "$OUTPUT_FILE"
                echo "$line" >> "$OUTPUT_FILE"
                echo "" >> "$OUTPUT_FILE"
                rate_drop_count=$((rate_drop_count + 1))
                problem_count=$((problem_count + 1))
            fi
        else
            connection_info="$line"
        fi
    done < "$temp_ss_output"

    # 如果没有速率问题，简洁输出
    if [ "$found_rate_drop" = false ]; then
        echo "速率问题：（pacing_rate:0或delivery_rate:0） - 未发现" >> "$OUTPUT_FILE"
    fi

    # 检查MTU问题
    echo "----------------------------------------" >> "$OUTPUT_FILE"
    found_mtu_problem=false
    while IFS= read -r line; do
        if [[ "$line" =~ ^[[:space:]] ]]; then
            if [[ "$line" =~ pmtu:([0-9]+) ]]; then
                pmtu=${BASH_REMATCH[1]}
                # 简化为整数比较
                if [ "$pmtu" -lt 1500 ]; then
                    if [ "$found_mtu_problem" = false ]; then
                        echo "MTU问题：（pmtu:<1500）" >> "$OUTPUT_FILE"
                        echo "" >> "$OUTPUT_FILE"
                        found_mtu_problem=true
                    fi
                    echo "$connection_info" >> "$OUTPUT_FILE"
                    echo "$line" >> "$OUTPUT_FILE"
                    echo "" >> "$OUTPUT_FILE"
                    mtu_problem_count=$((mtu_problem_count + 1))
                    problem_count=$((problem_count + 1))
                fi
            fi
        else
            connection_info="$line"
        fi
    done < "$temp_ss_output"

    # 如果没有MTU问题，简洁输出
    if [ "$found_mtu_problem" = false ]; then
        echo "MTU问题：（pmtu:<1500） - 未发现" >> "$OUTPUT_FILE"
    fi

    # 检查乱序包问题
    echo "----------------------------------------" >> "$OUTPUT_FILE"
    found_reordering=false
    while IFS= read -r line; do
        if [[ "$line" =~ ^[[:space:]] ]]; then
            if [[ "$line" =~ reordering ]]; then
                if [ "$found_reordering" = false ]; then
                    echo "乱序包：（存在reordering）" >> "$OUTPUT_FILE"
                    echo "" >> "$OUTPUT_FILE"
                    found_reordering=true
                fi
                echo "$connection_info" >> "$OUTPUT_FILE"
                echo "$line" >> "$OUTPUT_FILE"
                echo "" >> "$OUTPUT_FILE"
                reordering_count=$((reordering_count + 1))
                problem_count=$((problem_count + 1))
            fi
        else
            connection_info="$line"
        fi
    done < "$temp_ss_output"

    # 如果没有乱序包问题，简洁输出
    if [ "$found_reordering" = false ]; then
        echo "乱序包：（存在reordering） - 未发现" >> "$OUTPUT_FILE"
    fi

    # 问题总结
    echo "----------------------------------------" >> "$OUTPUT_FILE"
    if [ "$problem_count" -eq 0 ]; then
        echo "未发现任何TCP连接问题。" >> "$OUTPUT_FILE"
    else
        echo "问题TCP连接统计：" >> "$OUTPUT_FILE"
        echo "  高延迟连接数: $high_rtt_count" >> "$OUTPUT_FILE"
        echo "  重传连接数: $high_retrans_count" >> "$OUTPUT_FILE"
        echo "  窗口耗尽连接数: $window_exhausted_count" >> "$OUTPUT_FILE"
        echo "  速率问题连接数: $rate_drop_count" >> "$OUTPUT_FILE"
        echo "  MTU问题连接数: $mtu_problem_count" >> "$OUTPUT_FILE"
        echo "  乱序包连接数: $reordering_count" >> "$OUTPUT_FILE"
        echo "总计发现 $problem_count 个TCP连接问题。" >> "$OUTPUT_FILE"
    fi
}

# 初始化计数器
count=0
total_samples=$((DURATION > 0 ? DURATION / INTERVAL : 1))

# 绘制进度条函数
draw_progress_bar() {
    local percent=$1
    local width=50
    local completed=$((width * percent / 100))
    local remaining=$((width - completed))
    local bar=$(printf "%${completed}s" | tr " " "#")
    local space=$(printf "%${remaining}s")
    printf "\r${YELLOW}[${bar}${space}] %3d%%${NC}" $percent
}

echo -e "${BLUE}开始监控TCP连接性能...${NC}"
echo -e "监控网卡: ${INTERFACE}"
echo -e "收集频率: ${INTERVAL}秒"
echo -e "收集时长: ${DURATION}秒"
echo -e "输出文件: $OUTPUT_FILE"
echo -e "按 Ctrl+C 可以提前结束监控\n"

# 执行数据收集
if [ "$DURATION" -gt 0 ]; then
    # 多次收集模式
    start_time=$(date +%s)
    end_time=$((start_time + DURATION))
    current_time=$(date +%s)
    
    # 显示初始进度条
    draw_progress_bar 0
    
    while [ "$current_time" -lt "$end_time" ]; do
        collect_data
        
        count=$((count + 1))
        
        # 计算下一次采集前需要等待的时间
        sleep_time=$INTERVAL
        if [ "$((current_time + INTERVAL))" -gt "$end_time" ]; then
            # 确保不会超过总时长
            sleep_time=$((end_time - current_time))
        fi
        
        if [ "$sleep_time" -gt 0 ]; then
            sleep "$sleep_time"
        fi
        
        current_time=$(date +%s)
        
        # 更新同一行的进度条
        progress=$((100 * (current_time - start_time) / DURATION))
        draw_progress_bar $progress
    done
    
    # 完成时显示100%的进度条
    draw_progress_bar 100
    echo -e "\n\n${GREEN}正在生成报告...${NC}"
    
    # 添加多次收集的总结
    echo -e "\n\n" >> "$OUTPUT_FILE" # 增加额外的空行
    echo "================================================================" >> "$OUTPUT_FILE"
    echo "=== 多次采集总结 ===" >> "$OUTPUT_FILE"
    echo "================================================================" >> "$OUTPUT_FILE"
    echo -e "\n" >> "$OUTPUT_FILE" # 再增加一行空行
    echo "监控网卡: ${INTERFACE}" >> "$OUTPUT_FILE"
    echo "采集开始时间: $(date)" >> "$OUTPUT_FILE"
    echo "采集结束时间: $(date)" >> "$OUTPUT_FILE"
    echo "总采集时长: $DURATION 秒" >> "$OUTPUT_FILE"
    echo "采集频率: 每 $INTERVAL 秒" >> "$OUTPUT_FILE"
    echo "总采集次数: $count 次" >> "$OUTPUT_FILE"
    
    # 可以在这里添加多次采集数据的汇总分析
    # ...
    
else
    # 单次收集模式
    collect_data
fi

echo "----------------------------------------" >> "$OUTPUT_FILE"
echo "报告结束时间: $(date)" >> "$OUTPUT_FILE"

# 调用清理函数
cleanup
trap '' EXIT  # 清除EXIT陷阱，防止再次触发cleanup
exit 0 